from os import path
from collections import defaultdict
import numpy as np
from ppgnss import gnss_utils
import matplotlib.pyplot as plt

obj_dir = "../data/snr"
sites = ["B216", "B240", "B286", "PD02"]
for site in sites:
    obj_filename = path.join(obj_dir, "%s.obj"%site)
    data = gnss_utils.loadobject(obj_filename)
    xr_ele, xr_azi, xr_snrs = data["ele"], data["azi"], data["snr"]
    ele_min = 1.25
    azi_min = 2.5
    ele_step_size = 2.5
    azi_step_size = 5
    nrows = int(90/ele_step_size)
    ncols = int(360/azi_step_size)
    nsat = len(xr_ele.coords["prn"])
    invalid_idx = xr_ele < 0
    idx_row_grid = np.array(np.int32(np.floor(xr_ele/ele_step_size)))
    idx_col_grid = np.array(np.int32(np.floor(xr_azi/azi_step_size))) 
    xx_grid = azi_min + idx_col_grid * azi_step_size
    yy_grid = ele_min + idx_row_grid * ele_step_size
    eles = np.arange(ele_min, 90, ele_step_size)
    azis = np.arange(azi_min, 360, azi_step_size)
    idx_row_grid[invalid_idx] = -9999
    idx_col_grid[invalid_idx] = -9999
    for band in xr_snrs.coords["band"].values:
        strband = str(band)
        sattype = strband[:3]
        idx_gnss = xr_ele.coords["gnss"].values == sattype
        band_grid = [[[list() for icol in range(ncols)] for irow in range(nrows)] for isat in range(nsat)]
        count_grid = np.zeros((nsat, nrows, ncols))
        for iepoch, epoch_snr in enumerate(xr_snrs.loc[strband]):
            iirow = idx_row_grid[idx_gnss, iepoch].reshape(-1)
            iicol = idx_col_grid[idx_gnss, iepoch].reshape(-1)
            satlist = np.array(range(nsat))
            valid_idx_sat = iirow > 0
            for isat, irow, icol in zip(satlist[valid_idx_sat], iirow[valid_idx_sat], iicol[valid_idx_sat]):
                band_grid[isat][irow][icol].append(epoch_snr.coords["time"].values)
                count_grid[isat][irow][icol] += 1
                # print(isat, irow, icol, xr_ele[:, iepoch, isat].loc[sattype].values, xr_azi[:, iepoch, isat].loc[sattype].values)
        count_grid[count_grid == 0] = np.nan
        fig = plt.figure(figsize=(80/2.54, 80/2.54))
        fig, axes = plt.subplots(nrows=int(np.ceil(nsat/6)), ncols=6, subplot_kw={'projection': 'polar'})
        for isat in range(nsat):
            i = isat//6
            j = isat%6
            # print(isat, i, j)
            axes[i][j].pcolor(np.deg2rad(eles), np.deg2rad(azis), np.transpose(count_grid[isat])) 
            axes[i][j].set_theta_zero_location("N")
        plt.savefig("%s_%s.png" % (site, strband), dpi=300)
        plt.close()
        break
    break